import os
import h5py
import numpy as np
from pathlib import Path


def check_features_for_nan(h5_path):
    h5_path = Path(h5_path)
    if not h5_path.is_file():
        print(f"❌ 文件不存在: {h5_path}")
        return True

    try:
        with h5py.File(h5_path, "r") as f:
            if "features" not in f:
                print(f"⚠️ 警告: 文件 {h5_path} 中没有 'features' 数据集")
                return True

            dataset = f["features"]
            print(f"  Dtype: {dataset.dtype}")
            print(f"  Shape: {dataset.shape}")

            feats = dataset[:]  # 可能是 float16

            # 安全地计算统计量：转为 float32
            if feats.dtype == np.float16:
                feats_for_stats = feats.astype(np.float32)
            else:
                feats_for_stats = feats

            print(
                f"  Min: {np.nanmin(feats_for_stats):.4f}, Max: {np.nanmax(feats_for_stats):.4f}"
            )
            print(
                f"  Mean: {np.nanmean(feats_for_stats):.4f}, Std: {np.nanstd(feats_for_stats):.4f}"
            )

            # NaN/Inf 检查用原始数据（支持 float16）
            nan_count = np.isnan(feats).sum()
            inf_count = np.isinf(feats).sum()
            total = feats.size

            if nan_count > 0 or inf_count > 0:
                print(f"❗ NaN: {nan_count}, Inf: {inf_count} / {total}")
                return True
            else:
                print(f"✅ 无 NaN/Inf: {h5_path}")
                return False

    except Exception as e:
        print(f"💥 读取失败: {e}")
        return True


def load_prediction_from_h5(h5_path):
    """
    从指定的 .h5 文件中读取 'prediction' 数据集。

    参数:
        h5_path (str 或 Path): .h5 文件路径

    返回:
        np.ndarray 或 None: 成功时返回 prediction 数据的 NumPy 数组，失败时返回 None
    """
    h5_path = Path(h5_path)
    if not h5_path.is_file():
        print(f"错误: 文件不存在 - {h5_path}")
        return None

    try:
        with h5py.File(h5_path, "r") as f:
            if "prediction" not in f and "features" not in f:
                print(f"错误: 文件 {h5_path} 中没有 'prediction' or 'features' 数据集")
                print(f"可用的数据集: {list(f.keys())}")
                return None

            # dataset = f["prediction"]
            dataset = f["features"]
            
            data = np.array(dataset)
            print(f"成功加载: {h5_path} | Shape: {data.shape} | Dtype: {data.dtype}")
            return data

    except Exception as e:
        print(f"读取失败: {e}")
        return None


def process_h5_or_directory(input_path):
    """
    处理输入路径：如果是目录，则遍历其中的 .h5 文件；如果是 .h5 文件，则直接处理。

    参数:
        input_path (str 或 Path): 输入路径（可以是目录或 .h5 文件）
    """
    input_path = Path(input_path)

    if input_path.is_file() and input_path.suffix == ".h5":
        # 情况1: 直接是一个 .h5 文件
        print(f"输入路径是一个 .h5 文件，直接处理: {input_path}")
        data = load_prediction_from_h5(input_path)
        if data is not None:
            # 👉 在这里添加你自己的处理逻辑 👈
            pass

    elif input_path.is_dir():
        # 情况2: 是一个目录，遍历其中的 .h5 文件
        print(f"输入路径是一个目录，正在遍历 .h5 文件: {input_path}")
        pred_files = list(input_path.glob("*.h5"))
        if not pred_files:
            print(f"警告: 目录 {input_path} 中未找到任何 .h5 文件")
            return

        for h5_path in pred_files:
            print(f"处理文件: {h5_path}")
            data = load_prediction_from_h5(h5_path)
            if data is not None:
                # 👉 在这里添加你自己的处理逻辑 👈
                pass

    else:
        print(f"错误: 输入路径既不是 .h5 文件，也不是有效目录: {input_path}")


# ========================
# 使用示例
# ========================

# 你可以修改这一行为以下任意一种：
#   - 一个目录路径（批量处理）
#   - 一个具体的 .h5 文件路径（单文件处理）

input_path = Path(
    "/data0/lcy/data/LNM/LNM_Zhujiang_univ1_processed/20x_512px_0px_overlap/features_uni_v1"
)
for item in input_path.iterdir():
    if item.is_file() and item.suffix == ".h5":
        check_features_for_nan(item)
    elif item.is_dir():
        for h5_file in item.glob("*.h5"):
            check_features_for_nan(h5_file)

# 示例：也可以直接指定单个文件
# input_path = Path("/data0/lcy/data/LNM/.../case123.h5")

# 执行处理
# process_h5_or_directory(input_path)
